"""数据预处理模块：负责训练集和测试集的完整预处理流程"""

import ast
import pandas as pd
from typing import Optional, List, Dict, Any
from .config import (
    ENCODING,
    TRAIN_FIELD_MAPPING,
    TEST_FIELD_MAPPING,
    TEST_DROP_FIELDS
)
from .text_cleaner import TextCleaner


class DataPreprocessor:
    """数据预处理类：处理训练集和测试集的完整流程"""

    def __init__(self):
        """初始化预处理工具"""
        self.text_cleaner = TextCleaner()

    def load_train_data(self, file_path: str = 'data/train.csv') -> pd.DataFrame:
        """
        加载训练集数据

        Args:
            file_path: 训练集文件路径，默认在data目录下

        Returns:
            加载的训练集DataFrame
        """
        try:
            df = pd.read_csv(file_path, encoding=ENCODING)
            print(f"成功加载训练集，共{len(df)}条记录")
            return df
        except Exception as e:
            print(f"加载训练集失败: {str(e)}")
            raise

    def load_test_data(self, file_path: str = 'data/test.csv') -> pd.DataFrame:
        """
        加载测试集数据

        Args:
            file_path: 测试集文件路径，默认在data目录下

        Returns:
            加载的测试集DataFrame
        """
        try:
            df = pd.read_csv(file_path, encoding=ENCODING)
            print(f"成功加载测试集，共{len(df)}条记录")
            return df
        except Exception as e:
            print(f"加载测试集失败: {str(e)}")
            raise

    def standardize_train_format(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        标准化训练集格式

        Args:
            df: 原始训练集DataFrame

        Returns:
            格式标准化后的训练集DataFrame
        """
        # 重命名字段
        df_standardized = df.rename(columns=TRAIN_FIELD_MAPPING)

        # 处理答案字段，转换为列表格式
        df_standardized['correct_options'] = df_standardized['correct_options'].apply(
            lambda x: ast.literal_eval(x) if pd.notna(x) else []
        )

        return df_standardized

    def standardize_test_format(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        标准化测试集格式

        Args:
            df: 原始测试集DataFrame

        Returns:
            格式标准化后的测试集DataFrame
        """
        # 重命名字段
        df_standardized = df.rename(columns=TEST_FIELD_MAPPING)

        # 删除冗余字段
        df_standardized = df_standardized.drop(columns=TEST_DROP_FIELDS, errors='ignore')

        return df_standardized

    def clean_text_fields(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        清洗文本字段（问题和选项）

        Args:
            df: 标准化后的DataFrame

        Returns:
            文本清洗后的DataFrame
        """
        # 复制DataFrame避免修改原始数据
        df_clean = df.copy()

        # 清洗问题文本
        df_clean['cleaned_question'] = df_clean['raw_question'].apply(
            self.text_cleaner.clean_general_text
        )

        # 清洗选项文本
        def clean_options(options_str: Any) -> List[str]:
            if pd.isna(options_str):
                return []
            try:
                options_list = ast.literal_eval(str(options_str))
                if not isinstance(options_list, list):
                    return []
                return [self.text_cleaner.clean_option_text(opt) for opt in options_list]
            except:
                return []

        df_clean['cleaned_options_list'] = df_clean['raw_options'].apply(clean_options)

        return df_clean

    def split_options(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        将选项拆分为独立字段（opt_A~opt_E）

        Args:
            df: 清洗后的DataFrame

        Returns:
            选项拆分后的DataFrame
        """
        # 复制DataFrame避免修改原始数据
        df_split = df.copy()

        # 初始化选项字段
        opt_columns = ['opt_A', 'opt_B', 'opt_C', 'opt_D', 'opt_E']
        for col in opt_columns:
            df_split[col] = ''

        # 填充选项内容
        for idx, options_list in df_split['cleaned_options_list'].items():
            for i, opt_content in enumerate(options_list[:5]):  # 只处理前5个选项
                if i < len(opt_columns):
                    df_split.at[idx, opt_columns[i]] = opt_content

        # 添加选项计数字段
        df_split['option_count'] = df_split['cleaned_options_list'].apply(
            lambda x: min(len(x), 5)
        )

        # 删除临时字段
        df_split = df_split.drop(columns=['cleaned_options_list'], errors='ignore')

        return df_split

    def preprocess_train(self, file_path: str = 'data/train.csv') -> pd.DataFrame:
        """
        完整的训练集预处理流程

        Args:
            file_path: 训练集文件路径

        Returns:
            预处理完成的训练集DataFrame
        """
        print("开始处理训练集...")
        df = self.load_train_data(file_path)
        df = self.standardize_train_format(df)
        df = self.clean_text_fields(df)
        df = self.split_options(df)
        print("训练集处理完成")
        return df

    def preprocess_test(self, file_path: str = 'data/test.csv') -> pd.DataFrame:
        """
        完整的测试集预处理流程

        Args:
            file_path: 测试集文件路径

        Returns:
            预处理完成的测试集DataFrame
        """
        print("开始处理测试集...")
        df = self.load_test_data(file_path)
        df = self.standardize_test_format(df)
        df = self.clean_text_fields(df)
        df = self.split_options(df)
        print("测试集处理完成")
        return df

    def save_preprocessed_data(self, df: pd.DataFrame, output_path: str) -> None:
        """
        保存预处理后的数据

        Args:
            df: 预处理完成的DataFrame
            output_path: 输出文件路径
        """
        try:
            # 创建输出目录（如果不存在）
            import os
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)

            df.to_csv(output_path, encoding=ENCODING, index=False)
            print(f"预处理后的数据已保存至: {output_path}")
        except Exception as e:
            print(f"保存数据失败: {str(e)}")
            raise
